import os

from sklearn.model_selection import train_test_split

from option import parse_args


def save_file(root_dir, actors, file_name):
    with open(os.path.join(root_dir, file_name), 'w') as f:
        for video_name in actors:
            f.write(video_name + '\n')


def main():
    args = parse_args()
    assert os.path.exists(args.root_dir)
    actors = [f'{i:02d}' for i in range(1, 29)]
    train_actors, test_actors = train_test_split(actors, test_size=0.28, shuffle=False)
    val_actors, test_actors = train_test_split(test_actors, test_size=0.5, shuffle=False)
    print(train_actors, val_actors, test_actors)
    save_file(args.root_dir, train_actors, 'train.txt')
    save_file(args.root_dir, val_actors, 'val.txt')
    save_file(args.root_dir, test_actors, 'test.txt')


if __name__ == '__main__':
    # Total: 2000
    main()
